import numpy as np
import matplotlib.pyplot as plt
from matplotlib.colors import ListedColormap


def read_map(file_path):
    with open(file_path, 'r') as file:
        lines = file.readlines()
    map_data = [list(map(int, list(line.strip()))) for line in lines]
    return np.array(map_data)


actions = [(0, -1), (1, 0), (0, 1), (-1, 0)]


def create_environment(map_data):
    grid_size = len(map_data)
    num_states = grid_size * grid_size
    num_actions = len(actions)
    P = np.zeros((num_states, num_actions, num_states))
    R = np.full((num_states, num_actions), 0)
    start_state = None
    goal_state = None
    for i in range(grid_size):
        for j in range(grid_size):
            state = i * grid_size + j
            for action_idx, (di, dj) in enumerate(actions):
                ni, nj = i + di, j + dj
                if 0 <= ni < grid_size and 0 <= nj < grid_size:
                    next_state = ni * grid_size + nj
                    if map_data[ni, nj] != 1:
                        if map_data[ni, nj] == 3:
                            R[state, action_idx] = 1
                        P[state, action_idx, next_state] = 1
                    else:
                        P[state, action_idx, next_state] = 2
                        R[state, action_idx] = -2
                else:
                    R[state, action_idx] = -1
                    P[state, action_idx, state] = 1
            if map_data[i, j] == 2:
                start_state = state
            elif map_data[i, j] == 3:
                goal_state = state
                R[state, :] = 1
    return P, R, start_state, goal_state, grid_size


def plot_2d_map(map_data):
    fig, ax = plt.subplots()
    cmap = ListedColormap(['white', 'black', 'green', 'red'])
    ax.imshow(map_data, cmap=cmap, origin='upper', vmin=0, vmax=3, aspect='equal')
    ax.set_xlim(-0.5, map_data.shape[1] - 0.5)
    ax.set_ylim(map_data.shape[0] - 0.5, -0.5)
    ax.set_xlabel('X Axis')
    ax.set_ylabel('Y Axis')
    ax.set_xticks(np.arange(-0.5, map_data.shape[1], 1), minor=True)
    ax.set_yticks(np.arange(-0.5, map_data.shape[0], 1), minor=True)
    ax.grid(which='minor', color='gray', linestyle='-', linewidth=0.5)
    ax.set_xticks([])
    ax.set_yticks([])
    plt.show()


if __name__ == '__main__':
    map_data = read_map('map20*20.txt')
    P, R, start_state, goal_state, grid_size = create_environment(map_data)
    print(R)
    Q = np.zeros([grid_size * grid_size, 4])
    alpha = 0.1
    gamma = 0.95
    num_episodes = 1000
    rs = np.zeros([num_episodes])

    for i in range(num_episodes):
        r_sum_i = 0
        t = 0
        done = False
        s = start_state
        while not done:
            a = np.argmax(Q[s, :] + np.random.randn(1, 4) * (1. / (i / 10 + 1)))
            s1 = np.argmax(P[s, a, :])
            if np.max(P[s, a, :]) != 1 or s1 == goal_state:
                done = True
            r = R[s, a]
            print(f'    step: %d, action: %d, next_status: %d, r: %d, done: %s' % (t, a, s1, r, done))
            Q[s, a] = (1 - alpha) * Q[s, a] + alpha * (r + gamma * np.max(Q[s1, :]))
            r_sum_i += r * gamma ** t
            s = s1
            t += 1
        rs[i] = r_sum_i
        print(f'epoch: %d,  reward: %f' % (i, r_sum_i))
    print(Q)

    r_cumsum = np.cumsum(np.insert(rs, 0, 0))
    r_cumsum = (r_cumsum[50:] - r_cumsum[:-50]) / 50
    plt.plot(r_cumsum)
    plt.show()

    s = start_state
    path = []
    total_reward = 0.0
    while s != goal_state:
        a = np.argmax(Q[s, :])
        path.append(int(s))
        s1 = np.argmax(P[s, a, :])
        total_reward += R[s, a]
        s = s1
    print(f"Best Path: {path}")
    print(f"Total Reward: {total_reward}")
    input("Press Enter to exit...")
